import os
import pickle
import numpy as np
import json


if __name__ == "__main__":
    index = 531
    seed_path = "/home/ubuntu/Ascend/phase2_20221231/dataset/pytorch/seed" \
                         "/torch.argsort/20221229_212949/torch.argsort_seeds.npz"
    model_input_path = "/home/ubuntu/Ascend/phase2_20221231/dataset/pytorch/input" \
                         "/torch.argsort/20221229_212949/torch.argsort_inputs.json"
    source_result_path = "/home/ubuntu/Ascend/phase2_20221231/results/pt2cann/pytorch_gpu" \
                         "/torch.argsort/20221229_212949/torch.argsort_%d_results.npz" % index
    follow_result_path = "/home/ubuntu/Ascend/phase2_20221231/results/pt2cann/om" \
                         "/torch.argsort/20221229_212949/torch.argsort_%d_results.npz" % index
    om_path = "/home/ubuntu/Ascend/phase2_20221231/models/om_pytorch/" \
              "torch.argsort/20221229_212949/torch.argsort_%d.om" % index
    os.system("cp %s pt_results" % om_path)

    with open(model_input_path, "r") as f:
        model_inputs = json.load(f)[index]["inputs"]

    model_input_dict = {}

    with open("pt_results/TopKD.pickle", "wb+") as f:
        key_name = "torch.argsort_%d" % index
        seed_dict = np.load(seed_path)
        for model_input in model_inputs:
            for key in model_input.keys():
                val = model_input[key]
                model_input_dict[key] = seed_dict[val]
        print("model_input_dict:")
        print(model_input_dict)

        print(np.load(source_result_path)[key_name])
        print(np.load(follow_result_path)[key_name][0])

        results = [model_input_dict,
                     np.load(source_result_path)[key_name],
                     np.load(follow_result_path)[key_name][0]]

        pickle.dump(results, f)




